const std = @import("std");
const math = std.math;
const mem = std.mem;
const lbmath = @import("./matrix/math/lbmath.zig");

const SvdError = error{
    RowsLessThanCols,
    NoConvergence,
};

/// Computes sqrt(a^2 + b^2) without destructive underflow/overflow.
/// Ported from Numerical Recipes PYTHAG macro.
inline fn pythag(comptime T: type, a: T, b: T) T {
    const at: T = @abs(a);
    const bt: T = @abs(b);
    if (at > bt) {
        const ct: T = bt / at;
        return at * math.sqrt(1 + ct * ct);
    } else if (bt > 0) {
        const ct: T = at / bt;
        return bt * math.sqrt(1 + ct * ct);
    }
    return 0;
}

/// Swap two elements in a flat slice
inline fn swap(comptime T: type, values: []T, left: usize, right: usize) void {
    const tmp: T = values[left];
    values[left] = values[right];
    values[right] = tmp;
}

/// Swap two columns in a matrix represented as [][]T
fn swapMatrixColumns(comptime T: type, matrix: [][]T, col1: usize, col2: usize) void {
    for (matrix) |row| {
        const tmp: T = row[col1];
        row[col1] = row[col2];
        row[col2] = tmp;
    }
}

/// Swap column col1 and col2 in both U (m x n) and V (n x n)
fn swapMatrixColumnsBoth(comptime T: type, u: [][]T, v: [][]T, col1: usize, col2: usize) void {
    // swap in U
    for (u) |row| {
        const tmp: T = row[col1];
        row[col1] = row[col2];
        row[col2] = tmp;
    }
    // swap in V
    for (v) |row| {
        const tmp: T = row[col1];
        row[col1] = row[col2];
        row[col2] = tmp;
    }
}

/// QuickSort-like partition sorting (descending) for singular values w,
/// and simultaneously rearrange columns of U and V.
fn sortSingularValuesAndVectors(
    comptime T: type,
    w: []T, // length n
    u: [][]T, // m x n
    v: [][]T, // n x n
    left: usize,
    right: usize,
) void {
    if (left >= right) return;

    const pivot_value = w[left];
    var m_pivot: usize = left;

    for ((left + 1)..(right + 1)) |i| {
        if (w[i] > pivot_value) { // descending
            m_pivot += 1;
            swap(T, w, m_pivot, i);
            swapMatrixColumnsBoth(T, u, v, m_pivot, i);
        }
    }

    swap(T, w, left, m_pivot);
    swapMatrixColumnsBoth(T, u, v, left, m_pivot);

    if (m_pivot > left) {
        sortSingularValuesAndVectors(T, w, u, v, left, m_pivot - 1);
    }
    if (m_pivot < right) {
        sortSingularValuesAndVectors(T, w, u, v, m_pivot + 1, right);
    }
}
/// Golub–Reinsch SVD, 基本等价于 Numerical Recipes 的 svdcmp（0-based 索引版本）
pub fn dsvd(
    comptime T: type,
    allocator: std.mem.Allocator,
    a_in: [][]T, // m x n
    m: usize,
    n: usize,
    w_out: []T, // length n
    u_out: [][]T, // m x n, 输出 U
    v_out: [][]T, // n x n, 输出 V
) !void {
    if (m < n) return SvdError.RowsLessThanCols;

    // 拷贝 A -> U，后续所有操作在 u_out 上进行
    for (0..m) |i| {
        mem.copy(T, u_out[i][0..n], a_in[i][0..n]);
    }

    var rv1 = try allocator.alloc(T, n);
    defer allocator.free(rv1);

    var g: T = 0;
    var scale: T = 0;
    var anorm: T = 0;

    // -------------------------
    // Part 1: 双对角化（Householder）
    // -------------------------
    var i: usize = 0;
    while (i < n) : (i += 1) {
        const l = i + 1;

        // 先把上一轮右变换的信息存入 rv1[i]
        rv1[i] = scale * g;

        g = 0;
        var s: T = 0;
        scale = 0;

        // ---- 左侧变换（列 i）----
        if (i < m) {
            // 计算列 i 的范数缩放因子
            for (i..m) |k| {
                scale += @abs(u_out[k][i]);
            }

            if (scale != 0) {
                // 归一化该列并求平方和
                for (i..m) |k| {
                    u_out[k][i] /= scale;
                    s += u_out[k][i] * u_out[k][i];
                }
                const f = u_out[i][i];
                // g = -sign(f) * sqrt(s)
                g = if (f >= 0) -math.sqrt(s) else math.sqrt(s);
                const h = f * g - s;
                u_out[i][i] = f - g;

                if (l < n) {
                    // 用这个 Householder 向量更新右侧所有列 j >= l
                    for (l..n) |j| {
                        var sum: T = 0;
                        for (i..m) |k| {
                            sum += u_out[k][i] * u_out[k][j];
                        }
                        const ftmp = sum / h;
                        for (i..m) |k| {
                            u_out[k][j] += ftmp * u_out[k][i];
                        }
                    }
                }

                // 恢复列缩放
                for (i..m) |k| {
                    u_out[k][i] *= scale;
                }
            }
        }
        w_out[i] = scale * g;

        // ---- 右侧变换（行 i）----
        g = 0;
        s = 0;
        scale = 0;

        if (i < m and l < n) {
            // 计算行 i 从列 l..n-1 的缩放因子
            for (l..n) |k| {
                scale += @abs(u_out[i][k]);
            }

            if (scale != 0) {
                // 归一化该行的右半部分并求平方和
                for (l..n) |k| {
                    u_out[i][k] /= scale;
                    s += u_out[i][k] * u_out[i][k];
                }
                const f = u_out[i][l];
                g = if (f >= 0) -math.sqrt(s) else math.sqrt(s);
                const h = f * g - s;
                u_out[i][l] = f - g;

                // Householder 向量存入 rv1[k]
                for (l..n) |k| {
                    rv1[k] = u_out[i][k] / h;
                }

                // 用该 Householder 向量更新所有后续行 j >= i+1
                if (i + 1 < m) {
                    for (l..m) |j_row| {
                        var sum2: T = 0;
                        for (l..n) |k| {
                            sum2 += u_out[j_row][k] * u_out[i][k];
                        }
                        for (l..n) |k| {
                            u_out[j_row][k] += sum2 * rv1[k];
                        }
                    }
                }

                // 恢复行缩放
                for (l..n) |k| {
                    u_out[i][k] *= scale;
                }
            }
        }

        // anorm 用于 QR 迭代时的收敛判定
        const wi = @abs(w_out[i]);
        const ei = if (i < n) @abs(rv1[i]) else 0;
        anorm = if (wi + ei > anorm) wi + ei else anorm;
    }

    // -------------------------
    // Part 2: 累计右变换到 V（n x n）
    // -------------------------
    // 先把 V 初始化为单位阵
    for (0..n) |r| {
        for (0..n) |c| v_out[r][c] = 0;
        v_out[r][r] = 1;
    }

    g = 0;
    var l_idx: usize = n - 1; // 对应 NR 中的 l（1-based）转成 0-based

    var i_v: isize = @intCast(n - 1);

    while (i_v >= 0) : (i_v -= 1) {
        const izero: usize = @intCast(i_v);

        if (izero < n - 1) {
            if (g != 0) {
                // v[j][i] = (u[i][j] / u[i][l]) / g
                for (l_idx..n) |j| {
                    v_out[j][izero] = (u_out[izero][j] / u_out[izero][l_idx]) / g;
                }
                // 更新 V 的右侧块
                for (l_idx..n) |j| {
                    var sum: T = 0;
                    for (l_idx..n) |k| {
                        sum += u_out[izero][k] * v_out[k][j];
                    }
                    for (l_idx..n) |k| {
                        v_out[k][j] += sum * v_out[k][izero];
                    }
                }
            }
            // 清掉 i 行/列右边的部分
            for (l_idx..n) |j| {
                v_out[izero][j] = 0;
                v_out[j][izero] = 0;
            }
        }

        v_out[izero][izero] = 1;
        g = rv1[izero];
        l_idx = izero;
    }

    // -------------------------
    // Part 2: 累计左变换到 U（m x n）
    // -------------------------
    // 注意这里 U 存在 u_out 里
    const nm: usize = if (m < n) m else n;
    var i_u: isize = @intCast(nm - 1);
    while (i_u >= 0) : (i_u -= 1) {
        const izero: usize = @intCast(i_u);
        const l = izero + 1;
        g = w_out[izero];

        // 把当前行 izero 的右边部分清零
        if (l < n) {
            for (l..n) |j| {
                u_out[izero][j] = 0;
            }
        }

        if (g != 0) {
            g = 1 / g;
            if (l < n) {
                for (l..n) |j| {
                    var sum: T = 0;
                    // 根据 NR：k 从 l..m-1
                    for (l..m) |k| {
                        sum += u_out[k][izero] * u_out[k][j];
                    }
                    const f = (sum / u_out[izero][izero]) * g;
                    for (izero..m) |k| {
                        u_out[k][j] += f * u_out[k][izero];
                    }
                }
            }
            // 缩放列 izero
            for (izero..m) |k| {
                u_out[k][izero] *= g;
            }
        } else {
            // g == 0，则整列为 0
            for (izero..m) |k| {
                u_out[k][izero] = 0;
            }
        }
        u_out[izero][izero] += 1;
    }

    // -------------------------
    // Part 3: QR 迭代对双对角矩阵对角化
    // -------------------------
    var k: isize = @intCast(n - 1);
    while (k >= 0) : (k -= 1) {
        const k0: usize = @intCast(k);

        var its: usize = 0;
        while (its < 30) : (its += 1) {
            // 1. 查找拆分点 l
            var l: isize = k;
            var l0: usize = @intCast(l);
            var flag = true;

            while (l >= 0) : (l -= 1) {
                l0 = @intCast(l);
                if (l0 == 0) break;
                const nm0 = l0 - 1;
                if (@abs(rv1[l0]) + anorm == anorm) {
                    flag = false;
                    break;
                }
                if (@abs(w_out[nm0]) + anorm == anorm) {
                    break;
                }
            }

            if (flag) {
                // 2. 消去 rv1[l]
                var c: T = 0;
                var s2: T = 1;
                var iqr: isize = l;
                while (iqr <= k) : (iqr += 1) {
                    const i_idx: usize = @intCast(iqr);
                    const nm0: isize = @intCast(iqr - 1);
                    const f = s2 * rv1[i_idx];
                    rv1[i_idx] = c * rv1[i_idx];
                    if (@abs(f) + anorm != anorm) {
                        const g2 = w_out[i_idx];
                        var h = pythag(T, f, g2);
                        w_out[i_idx] = h;
                        h = 1 / h;
                        c = g2 * h;
                        s2 = -f * h;
                        // 旋转 U 的两列 (nm0, i_idx)
                        for (0..m) |j| {
                            const y = u_out[j][nm0];
                            const z = u_out[j][i_idx];
                            u_out[j][nm0] = y * c + z * s2;
                            u_out[j][i_idx] = z * c - y * s2;
                        }
                    }
                }
            }

            const z = w_out[k0];
            // 3. 判断是否收敛
            if (l == k) {
                if (z < 0) {
                    w_out[k0] = -z;
                    // 对应列的 V 乘以 -1
                    for (0..n) |j| {
                        v_out[j][k0] = -v_out[j][k0];
                    }
                }
                break; // 该 k 收敛
            }

            if (its == 29) return SvdError.NoConvergence;

            // 4. Wilkinson shift
            const l_idx2: usize = @intCast(l);
            const nm0 = k0 - 1;
            var x = w_out[l_idx2];
            var y = w_out[nm0];
            var g2 = rv1[nm0];
            var h2 = rv1[k0];
            var f = ((y - z) * (y + z) + (g2 - h2) * (g2 + h2)) / (2 * h2 * y);
            g2 = pythag(T, f, 1);
            f = ((x - z) * (x + z) + h2 * (y / (f + (if (f >= 0) g2 else -g2)) - h2)) / x;

            // 5. bulge chasing
            var c2: T = 1;
            var s3: T = 1;
            // for (var j_it: isize = @intCast(l_idx2); j_it <= @intCast(nm0); j_it += 1)
            for (l_idx2..(nm0 + 1)) |j_it| {
                const j0: usize = @intCast(j_it);
                const ione: usize = j0 + 1;

                var g3 = rv1[ione];
                y = w_out[ione];
                h2 = s3 * g3;
                g3 = c2 * g3;
                var z2 = pythag(T, f, h2);
                rv1[j0] = z2;
                c2 = f / z2;
                s3 = h2 / z2;
                f = x * c2 + g3 * s3;
                g3 = -x * s3 + g3 * c2;
                h2 = y * s3;
                y *= c2;

                // 作用到 V 的列 j0 和 ione
                for (0..n) |jj| {
                    const vx = v_out[jj][j0];
                    const vz = v_out[jj][ione];
                    v_out[jj][j0] = vx * c2 + vz * s3;
                    v_out[jj][ione] = vz * c2 - vx * s3;
                }

                z2 = pythag(T, f, h2);
                w_out[j0] = z2;
                if (z2 != 0) {
                    z2 = 1 / z2;
                    c2 = f * z2;
                    s3 = h2 * z2;
                }
                f = c2 * g3 + s3 * y;
                x = c2 * y - s3 * g3;

                // 作用到 U 的列 j0 和 ione
                for (0..m) |jj| {
                    const uy = u_out[jj][j0];
                    const uz = u_out[jj][ione];
                    u_out[jj][j0] = uy * c2 + uz * s3;
                    u_out[jj][ione] = uz * c2 - uy * s3;
                }
            }

            rv1[@intCast(l_idx2)] = 0;
            rv1[k0] = f;
            w_out[k0] = x;
        } // its
    } // k
    // At this point we have singular values in w_out (nonnegative), U in u_out (columns), V in v_out.
    // Sort singular values in descending order and rearrange U and V accordingly.
    if (n > 1) {
        sortSingularValuesAndVectors(T, w_out, u_out, v_out, 0, n - 1);
    }

    return;
}

// --- Test block (kept for completeness and verification) ---
// test "dsvd simple 2x2 matrix" {
//     const default_allocator = std.testing.allocator;

//     var a_data: [2][2]f64 = .{
//         .{1.0, 1.0},
//         .{1.0, 1.0},
//     };
//     var a_in: [][]f64 = &.{a_data[0][0..], a_data[1][0..]};

//     var w_out_data: [2]f64 = undefined;
//     var w_out: []f64 = &w_out_data;

//     var u_out_data: [2][2]f64 = undefined;
//     var u_out: [][]f64 = &.{u_out_data[0][0..], u_out_data[1][0..]};

//     var v_out_data: [2][2]f64 = undefined;
//     var v_out: [][]f64 = &.{v_out_data[0][0..], v_out_data[1][0..]};

//     try dsvd(default_allocator, a_in, 2, 2, w_out, u_out, v_out);

//     std.debug.print("SVD of A:\n", .{});
//     std.debug.print("  A = {any}\n", .{a_in});
//     std.debug.print("  W = {any}\n", .{w_out});
//     std.debug.print("  U = {any}\n", .{u_out});
//     std.debug.print("  V = {any}\n", .{v_out});

//     // Check singular values
//     try std.testing.expectApproxEqAbs(2.0, w_out[0], 1e-9);
//     try std.testing.expectApproxEqAbs(0.0, w_out[1], 1e-9);

//     const sqrt2_inv = 1.0 / math.sqrt(2.0);

//     // Check U (signs might vary)
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(u_out[0][0]), 1e-9);
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(u_out[1][0]), 1e-9);
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(u_out[0][1]), 1e-9);
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(u_out[1][1]), 1e-9);
//     // U columns orthogonal
//     try std.testing.expectApproxEqAbs(0.0, u_out[0][0] * u_out[0][1] + u_out[1][0] * u_out[1][1], 1e-9);

//     // Check V (signs might vary)
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(v_out[0][0]), 1e-9);
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(v_out[1][0]), 1e-9);
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(v_out[0][1]), 1e-9);
//     try std.testing.expectApproxEqAbs(sqrt2_inv, @abs(v_out[1][1]), 1e-9);
//     // V columns orthogonal
//     try std.testing.expectApproxEqAbs(0.0, v_out[0][0] * v_out[0][1] + v_out[1][0] * v_out[1][1], 1e-9);

//     // Test a non-square matrix (m > n)
//     // A = [[1, 2], [3, 4], [5, 6]] (3x2)
//     var a_tall_data: [3][2]f64 = .{
//         .{1.0, 2.0},
//         .{3.0, 4.0},
//         .{5.0, 6.0},
//     };
//     var a_tall_in: [][]f64 = &.{a_tall_data[0][0..], a_tall_data[1][0..], a_tall_data[2][0..]};

//     var w_tall_out_data: [2]f64 = undefined; // n=2 singular values
//     var w_tall_out: []f64 = &w_tall_out_data;

//     var u_tall_out_data: [3][2]f64 = undefined; // m x n = 3 x 2 for U
//     var u_tall_out: [][]f64 = &.{u_tall_out_data[0][0..], u_tall_out_data[1][0..], u_tall_out_data[2][0..]};

//     var v_tall_out_data: [2][2]f64 = undefined; // n x n = 2 x 2 for V
//     var v_tall_out: [][]f64 = &.{v_tall_out_data[0][0..], v_tall_out_data[1][0..]};

//     try dsvd(default_allocator, a_tall_in, 3, 2, w_tall_out, u_tall_out, v_tall_out);

//     std.debug.print("\nSVD of A (tall):\n", .{});
//     std.debug.print("  A = {any}\n", .{a_tall_in});
//     std.debug.print("  W = {any}\n", .{w_tall_out});
//     std.debug.print("  U = {any}\n", .{u_tall_out});
//     std.debug.print("  V = {any}\n", .{v_tall_out});

//     // Known singular values for [[1,2],[3,4],[5,6]] from NumPy
//     // W = [9.52551809, 0.51430058] (approx)
//     try std.testing.expectApproxEqAbs(9.52551809, w_tall_out[0], 1e-5);
//     try std.testing.expectApproxEqAbs(0.51430058, w_tall_out[1], 1e-5);

//     // Verification of A = U @ S @ V.T
//     // S_mat is a m x n matrix with w_out on diagonal
//     var s_mat_data: [3][2]f64 = .{
//         .{0.0, 0.0},
//         .{0.0, 0.0},
//         .{0.0, 0.0},
//     };
//     s_mat_data[0][0] = w_tall_out[0];
//     s_mat_data[1][1] = w_tall_out[1];
//     var s_mat: [][]f64 = &.{s_mat_data[0][0..], s_mat_data[1][0..], s_mat_data[2][0..]};

//     var us_prod_data: [3][2]f64 = undefined;
//     var us_prod: [][]f64 = &.{us_prod_data[0][0..], us_prod_data[1][0..], us_prod_data[2][0..]};

//     // Calculate U @ S
//     for (0..3) |i| {
//         for (0..2) |j| {
//             us_prod[i][j] = 0.0;
//             for (0..2) |k| { // U is m x n (3x2), S is n x n (2x2). Result is m x n (3x2)
//                 us_prod[i][j] += u_tall_out[i][k] * s_mat[k][j];
//             }
//         }
//     }

//     var usv_prod_data: [3][2]f64 = undefined;
//     var usv_prod: [][]f64 = &.{usv_prod_data[0][0..], usv_prod_data[1][0..], usv_prod_data[2][0..]};

//     // Calculate (U @ S) @ V.T
//     for (0..3) |i| {
//         for (0..2) |j| {
//             usv_prod[i][j] = 0.0;
//             for (0..2) |k| { // US is m x n (3x2), V.T is n x n (2x2). Result is m x n (3x2)
//                 usv_prod[i][j] += us_prod[i][k] * v_tall_out[j][k]; // v_tall_out[j][k] is element (k,j) of V, which is (j,k) of V.T
//             }
//         }
//     }

//     std.debug.print("  U @ S @ V.T = {any}\n", .{usv_prod});

//     // Compare A with U @ S @ V.T
//     for (0..3) |i| {
//         for (0..2) |j| {
//             try std.testing.expectApproxEqAbs(a_tall_in[i][j], usv_prod[i][j], 1e-5);
//         }
//     }
// }
